"""
K线订阅管理器

提供K线数据订阅、历史K线加载、实时K线更新处理等功能
可重用于不同的策略实现中
"""

import asyncio
import logging
import queue
import threading
import json
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional, Callable

from services.exchange.adapters.gate_adapter import GateAdapter
from services.exchange import KLineInterval

logger = logging.getLogger(__name__)

class KlineSubscriber:
    """K线订阅管理器，处理K线数据的订阅、加载和更新"""
    
    def __init__(self, symbol: str, exchange_adapter=None, use_testnet: bool = False, logger_name: str = None):
        """
        初始化K线订阅管理器
        
        参数:
            symbol: 交易对符号
            exchange_adapter: 交易所适配器，如果为None则自动创建GateAdapter
            use_testnet: 是否使用测试网
            logger_name: 日志记录器名称，如果为None则使用当前模块的logger
        """
        self.symbol = symbol
        
        # 设置日志记录器
        self.logger = logger_name if isinstance(logger_name, logging.Logger) else logging.getLogger(logger_name or __name__)
        
        # 初始化交易所适配器
        self.exchange = exchange_adapter or GateAdapter(use_testnet=use_testnet)
        
        # 初始化K线数据容器（保存100根K线）
        self.klines_data = []
        self.klines_lock = threading.Lock()  # 用于线程安全访问K线数据
        
        # 消息队列用于存储K线更新事件
        self.event_queue = queue.Queue()
        
        # 记录最后更新时间
        self.last_update_time = None
        
        # 回调函数，当k线更新时会调用此函数
        self.on_kline_update_callback = None
        
        # 记录K线内交易状态，避免同一K线内重复交易
        self.current_candle_time = None
        
    async def initialize(self, limit: int = 100):
        """
        初始化订阅器，加载历史K线数据
        
        参数:
            limit: 加载的历史K线数量
        
        返回:
            是否初始化成功
        """
        try:
            # 从交易所获取历史K线
            self.logger.info(f"正在获取 {self.symbol} 的历史K线数据...")
            klines = await self.exchange.get_klines(
                symbol=self.symbol,
                interval=KLineInterval.MINUTE_5,  # 使用5分钟K线
                limit=limit
            )
            
            if not klines:
                self.logger.error(f"未能获取 {self.symbol} 的历史K线数据")
                return False
            
            # 将K线数据保存
            with self.klines_lock:
                self.klines_data = klines
                df = self.convert_klines_to_dataframe(self.klines_data)
                
                # 打印最新的2根K线数据
                self.print_latest_klines(df, n=2)
            
            self.last_update_time = datetime.now()
            return True
        except Exception as e:
            self.logger.error(f"初始化 {self.symbol} K线订阅器时出错: {e}", exc_info=True)
            return False
    
    def convert_klines_to_dataframe(self, klines: List[Dict[str, Any]]) -> Any:
        """
        将K线数据列表转换为pandas DataFrame
        
        参数:
            klines: K线数据列表
            
        返回:
            pandas DataFrame
        """
        import pandas as pd
        df = pd.DataFrame([
            {
                'datetime': kline['datetime'],
                'open': kline['open'],
                'high': kline['high'],
                'low': kline['low'],
                'close': kline['close'],
                'volume': kline['volume']
            } for kline in klines
        ])
        
        # 确保按时间排序
        if not df.empty:
            df = df.sort_values('datetime').reset_index(drop=True)
        
        return df
    
    def print_latest_klines(self, df: Any, n: int = 2):
        """
        打印最新的n根K线数据
        
        参数:
            df: K线数据DataFrame
            n: 打印的K线数量
        """
        if df.empty or len(df) < n:
            self.logger.warning(f"K线数据不足 {n} 根，无法打印")
            return
        
        latest_klines = df.tail(n)
        self.logger.info(f"最新的 {n} 根K线数据 ({self.symbol}):")
        for idx, row in latest_klines.iterrows():
            self.logger.info(f"  时间: {row['datetime'].strftime('%Y-%m-%d %H:%M:%S')}, "
                       f"开: {row['open']:.6f}, "
                       f"高: {row['high']:.6f}, "
                       f"低: {row['low']:.6f}, "
                       f"收: {row['close']:.6f}, "
                       f"量: {row['volume']:.6f}")
    
    def handle_kline_update(self, data: List[Dict[str, Any]]):
        """
        处理K线数据更新
        
        参数:
            data: K线数据
        """
        self.logger.info(f"收到K线更新: {json.dumps(data, ensure_ascii=False)}")
        
        # 将更新事件放入队列
        self.event_queue.put(data)

    async def process_kline_updates(self):
        """处理K线更新事件的主循环"""
        while True:
            try:
                # 从队列中获取更新事件
                if self.event_queue.empty():
                    # 即使没有新数据，也定期进行回调
                    now = datetime.now()
                    if self.last_update_time is None or (now - self.last_update_time).total_seconds() > 5:
                        if len(self.klines_data) > 0 and self.on_kline_update_callback:
                            with self.klines_lock:
                                df = self.convert_klines_to_dataframe(self.klines_data)
                                await self.on_kline_update_callback(df)
                            self.last_update_time = now
                    
                    await asyncio.sleep(0.1)
                    continue
                
                data = self.event_queue.get()
                self.logger.info(f"处理K线更新事件")
                
                # 处理K线数据
                for kline in data:
                    contract = kline.get('n', '')  # 合约名称
                    # 确保是目标交易对
                    symbol = contract.replace('5m_', '')
                    if symbol != self.symbol:
                        self.logger.debug(f"跳过非目标交易对数据: {contract}")
                        continue
                    
                    # 解析K线数据
                    timestamp = kline.get('t', 0)
                    open_price = float(kline.get('o', 0))
                    high_price = float(kline.get('h', 0))
                    low_price = float(kline.get('l', 0))
                    close_price = float(kline.get('c', 0))
                    volume = float(kline.get('v', 0))
                    
                    # 将时间戳转换为北京时间
                    datetime_utc = datetime.utcfromtimestamp(timestamp)
                    # 北京时间比UTC早8小时
                    beijing_time = datetime_utc + timedelta(hours=8)
                    
                    self.logger.info(f"处理K线: 时间={beijing_time}, 交易对={contract}, 收盘价={close_price}")
                    
                    # 构建K线数据
                    new_kline = {
                        'datetime': beijing_time,
                        'open': open_price,
                        'high': high_price,
                        'low': low_price,
                        'close': close_price,
                        'volume': volume
                    }
                    
                    # 获取当前K线时间
                    current_candle_time = beijing_time.strftime('%Y-%m-%d %H:%M:%S')
                    self.current_candle_time = current_candle_time
                    
                    with self.klines_lock:
                        # 检查是否已有相同时间的K线（更新而不是添加）
                        updated = False
                        for i, k in enumerate(self.klines_data):
                            if k['datetime'].strftime('%Y-%m-%d %H:%M') == beijing_time.strftime('%Y-%m-%d %H:%M'):
                                self.logger.info(f"更新已存在的K线: {beijing_time.strftime('%Y-%m-%d %H:%M')}")
                                self.klines_data[i] = new_kline
                                updated = True
                                break
                        
                        # 如果是新的K线，添加到列表并保持100根K线
                        if not updated:
                            self.logger.info(f"添加新K线: {beijing_time.strftime('%Y-%m-%d %H:%M')}")
                            self.klines_data.append(new_kline)
                            # 保持最多100根K线
                            if len(self.klines_data) > 100:
                                self.klines_data = self.klines_data[-100:]
                        
                        # 转换为DataFrame并调用回调函数
                        df = self.convert_klines_to_dataframe(self.klines_data)
                        self.print_latest_klines(df, n=2)
                        
                        # 如果有回调函数，则调用它
                        if self.on_kline_update_callback:
                            await self.on_kline_update_callback(df)
                        
                        self.last_update_time = datetime.now()
                
                # 标记任务完成
                self.event_queue.task_done()
            
            except Exception as e:
                self.logger.error(f"处理K线更新时出错: {e}", exc_info=True)
                await asyncio.sleep(1)  # 出错时暂停一下

    async def subscribe(self, on_kline_update: Optional[Callable] = None):
        """
        订阅K线数据
        
        参数:
            on_kline_update: K线更新时的回调函数
            
        返回:
            是否成功订阅
        """
        # 设置回调函数
        self.on_kline_update_callback = on_kline_update
        
        # 修改为使用下划线分隔的交易对格式
        symbol_for_exchange = self.symbol.replace('-', '_')
        
        # 增加WebSocket连接异常处理
        retry_count = 0
        max_retries = 5
        
        while retry_count < max_retries:
            try:
                # 订阅前先关闭所有WebSocket连接，避免"socket is already opened"错误
                try:
                    self.exchange.close_all_ws_connections()
                    # 等待连接完全关闭
                    await asyncio.sleep(2)
                except Exception as e:
                    self.logger.warning(f"关闭WebSocket连接时出错: {e}, 继续尝试订阅...")
                
                # 订阅K线频道
                self.logger.info(f"订阅 {self.symbol} 的5分钟K线数据 (尝试 {retry_count+1}/{max_retries})...")
                if self.exchange.subscribe_kline(symbol_for_exchange, "5m", self.handle_kline_update):
                    self.logger.info(f"成功订阅 {self.symbol} K线数据")
                    
                    # 启动处理更新事件的任务和WebSocket健康检查
                    process_task = asyncio.create_task(self.process_kline_updates())
                    health_task = asyncio.create_task(self._websocket_health_check())
                    ping_task = asyncio.create_task(self._websocket_ping())
                    
                    return True
                else:
                    self.logger.warning(f"订阅 {self.symbol} K线失败，5秒后重试...")
                    retry_count += 1
                    await asyncio.sleep(5)
            except Exception as e:
                self.logger.error(f"订阅K线时出错: {str(e)}, 5秒后重试...")
                retry_count += 1
                await asyncio.sleep(5)
        
        self.logger.error(f"订阅 {self.symbol} K线失败，放弃重试")
        return False

    async def _websocket_ping(self):
        """每30秒发送一次pingpong消息，保持WebSocket连接活跃"""
        while True:
            try:
                self.logger.info(f"发送WebSocket ping消息...")
                self.exchange.send_ping()
                await asyncio.sleep(30)  # 每30秒执行一次
            except Exception as e:
                self.logger.info(f"发送ping消息出错: {e}", exc_info=True)
                await asyncio.sleep(5)  # 出错时暂停5秒

    async def _websocket_health_check(self):
        """检查WebSocket连接健康状态并在需要时重新连接"""
        reconnect_interval = 60  # 初始重连间隔60秒
        max_reconnect_interval = 300  # 最大重连间隔5分钟
        
        while True:
            try:
                # 检查WebSocket连接状态
                await asyncio.sleep(reconnect_interval)
                
                # 检查是否长时间没有收到K线更新
                current_time = datetime.now()
                
                # 检查最后一次K线更新时间
                if len(self.klines_data) > 0:
                    last_kline_time = self.klines_data[-1]['datetime']
                    time_since_last_update = (current_time - last_kline_time).total_seconds() / 60  # 转换为分钟
                    
                    # 如果超过15分钟没有更新，尝试重新订阅
                    if time_since_last_update > 15:
                        self.logger.warning(f"已超过 {time_since_last_update:.1f} 分钟未收到K线更新，尝试重新订阅")
                        
                        # 关闭现有连接
                        try:
                            self.exchange.close_all_ws_connections()
                            await asyncio.sleep(2)  # 等待连接完全关闭
                        except Exception as e:
                            self.logger.error(f"关闭WebSocket连接时出错: {e}")
                        
                        # 重新订阅
                        symbol_for_exchange = self.symbol.replace('-', '_')
                        subscription_success = self.exchange.subscribe_kline(symbol_for_exchange, "5m", self.handle_kline_update)
                        
                        if subscription_success:
                            self.logger.info(f"成功重新订阅 {self.symbol} K线数据")
                            # 重置重连间隔
                            reconnect_interval = 60
                        else:
                            self.logger.error(f"重新订阅 {self.symbol} K线失败")
                            # 增加重连间隔（指数退避）
                            reconnect_interval = min(reconnect_interval * 2, max_reconnect_interval)
                    else:
                        self.logger.debug(f"WebSocket连接状态正常，最近一次更新在 {time_since_last_update:.1f} 分钟前")
            
            except Exception as e:
                self.logger.error(f"WebSocket健康检查出错: {e}", exc_info=True)
                # 出错时增加重连间隔（指数退避）
                reconnect_interval = min(reconnect_interval * 2, max_reconnect_interval)
                await asyncio.sleep(5)
                
    def get_klines_data(self):
        """获取当前K线数据"""
        with self.klines_lock:
            return self.klines_data.copy()
    
    def get_current_price(self):
        """获取最新价格"""
        with self.klines_lock:
            if self.klines_data:
                return self.klines_data[-1]['close']
            return None 